import cv2
import numpy as np
from pathlib import Path

# COCO 17 keypoints骨架连线（每对为一条骨架）
SKELETON = [
    (0, 1), (0, 2), (1, 3), (2, 4),
    (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),
    (5, 11), (6, 12), (11, 12), (11, 13), (13, 15), (12, 14), (14, 16)
]

# 颜色
COLORS = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(255,0,255),(0,255,255),(128,128,0),(128,0,128),(0,128,128),(128,128,255),(255,128,128),(128,255,128),(128,128,128),(0,0,0),(255,255,255),(100,100,255),(255,100,100)]

# 可视化函数
def visualize_pose(image_path, label_path, output_path):
    img = cv2.imread(str(image_path))
    h, w = img.shape[:2]
    with open(label_path, 'r') as f:
        for line in f:
            items = line.strip().split()
            # 跳过非person类别
            if items[0] != '0':
                continue
            # 解析关键点
            kpts = []
            for i in range(17):
                x = float(items[1 + i*3]) * w
                y = float(items[2 + i*3]) * h
                v = float(items[3 + i*3])
                kpts.append((x, y, v))
            # 画点
            for idx, (x, y, v) in enumerate(kpts):
                if v > 0:
                    cv2.circle(img, (int(x), int(y)), 4, COLORS[idx%len(COLORS)], -1)
            # 画骨架
            for i, j in SKELETON:
                if kpts[i][2] > 0 and kpts[j][2] > 0:
                    cv2.line(img, (int(kpts[i][0]), int(kpts[i][1])), (int(kpts[j][0]), int(kpts[j][1])), (0,255,255), 2)
    cv2.imwrite(str(output_path), img)
    print(f'可视化结果已保存到: {output_path}')

if __name__ == '__main__':
    # 示例图片和标注
    image_path = Path('dataset/images/train/frame_045780.jpg')
    label_path = Path('dataset/labels/train/frame_045780.txt')
    output_path = Path('frame_045780_vis.jpg')
    visualize_pose(image_path, label_path, output_path) 